package cmd

import (
	"context"
	"fmt"
	"io/ioutil"
	"strconv"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/s3"
	"github.com/spf13/cobra"
	"github.com/tikv/client-go/txnkv"
)

var restoreCmd = &cobra.Command{
	Use:   "restore",
	Short: "Restore He3DB",
	Long:  "Welcome to use hbr for He3DB restore",
	Run:   runRestore,
}

func init() {
	rootCmd.AddCommand(restoreCmd)
}

func runRestore(cmd *cobra.Command, args []string) {
	access_key, _ := cmd.Flags().GetString("access_key")
	secret_key, _ := cmd.Flags().GetString("secret_key")
	endpoint, _ := cmd.Flags().GetString("endpoint")
	region, _ := cmd.Flags().GetString("region")
	bucket, _ := cmd.Flags().GetString("bucket")
	pd, _ := cmd.Flags().GetString("pd")
	backup_name, _ := cmd.Flags().GetString("name")

	if access_key == "" || secret_key == "" || endpoint == "" || region == "" || bucket == "" || pd == "" || backup_name == "" {
		fmt.Printf("PARAMETER ERROR!\n")
		return
	}

	client, err := txnkv.NewClient([]string{pd})
	if err != nil {
		fmt.Printf("Connect Tikv Error!\n%v\n", err)
		return
	}
	defer func() {
		client.Close()
	}()

	sess, err := session.NewSession(&aws.Config{
		Region:      aws.String(region),
		Endpoint:    aws.String(endpoint),
		Credentials: credentials.NewStaticCredentials(access_key, secret_key, ""),
		S3ForcePathStyle: aws.Bool(true),
	})
	if err != nil {
		fmt.Printf("Connect S3 Error!\n%v\n", err)
		return
	}
	s3_client := s3.New(sess)

	count := 0
	input := &s3.ListObjectsInput{
		Bucket: aws.String(bucket),
		Prefix: aws.String(backup_name),
	}
	for {
		resp, err := s3_client.ListObjects(input)
		if err != nil {
			fmt.Printf("S3 ListObjects Error!\n%v\n", err)
			return
		}

		for _, keys := range resp.Contents {
			out, err := s3_client.GetObject(&s3.GetObjectInput{
				Bucket: aws.String(bucket),
				Key:    aws.String(*keys.Key),
			})
			if err != nil {
				fmt.Printf("S3 GetObject Error!\n%v\n", err)
				return
			}
			defer out.Body.Close()

			data, err := ioutil.ReadAll(out.Body)
			if err != nil {
				fmt.Printf("out.Body.Read!\n%v\n", err)
				return
			}

			txn, err := client.Begin()
			if err != nil {
				fmt.Printf("Tikv Transaction Begin Error!\n%v\n", err)
				return
			}

			fmt.Printf("filename:%s\n", (*keys.Key)[len(backup_name)+1:])

			ret := make([]byte, (len(*keys.Key)-len(backup_name)-1)/2)
			index := 0
			for i := len(backup_name) + 1; i < len(*keys.Key); i += 2 {
				value, _ := strconv.ParseUint((*keys.Key)[i:i+2], 16, 8)
				ret[index] = byte(0xff & value)
				index++

			}

			if err := txn.Set(ret, data); err != nil {
				fmt.Printf("Tikv Set Error!\n%v\n", err)
				return
			}

			if err := txn.Commit(context.TODO()); err != nil {
				fmt.Printf("Tikv Transaction Commit Error!\n%v\n", err)
				return
			}
			count++
		}
		if resp.NextMarker == nil {
			fmt.Printf("Done!\n")
			break
		}
		input.Marker = resp.NextMarker
	}
	fmt.Printf("N:%v\n", count)
	fmt.Printf("Done!\n")
}
